import cv2
import numpy as np
from moviepy.editor import VideoFileClip
from pathlib import Path


class WatermarkRemover:
    """图片和视频去水印工具"""
    
    def __init__(self):
        self.supported_image_formats = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
        self.supported_video_formats = {'.mp4', '.avi', '.mov', '.mkv'}

    def remove_watermark_image(self, image_path, mask_coords, output_path=None):
        """
        图片去水印
        
        Args:
            image_path: 输入图片路径
            mask_coords: 水印区域坐标 [(x1,y1), (x2,y2)]
            output_path: 输出路径（可选）
        """
        # 读取图片
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError("无法读取图片")
        
        # 创建mask
        mask = np.zeros(img.shape[:2], dtype=np.uint8)
        x1, y1 = mask_coords[0]
        x2, y2 = mask_coords[1]
        mask[y1:y2, x1:x2] = 255
        
        # 使用OpenCV的inpainting
        result = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)
        
        # 保存结果
        if output_path is None:
            output_path = str(Path(image_path).with_suffix('')) + '_cleaned.jpg'
        
        cv2.imwrite(output_path, result)
        return output_path
    
    def detect_watermark_region(self, image_path, template_path=None):
        """
        自动检测水印区域（基于模板匹配）
        
        Args:
            image_path: 输入图片路径
            template_path: 水印模板路径（可选）
        """
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError("无法读取图片")
        
        # 如果没有模板，使用边缘检测找矩形区域
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 50, 150)
        
        # 找轮廓
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 筛选可能的矩形区域（水印通常是小矩形）
        watermark_regions = []
        for contour in contours:
            area = cv2.contourArea(contour)
            if 1000 < area < 50000:  # 根据实际调整
                x, y, w, h = cv2.boundingRect(contour)
                if w > 20 and h > 20:  # 最小尺寸限制
                    watermark_regions.append([(x, y), (x+w, y+h)])
        
        return watermark_regions
    
    def remove_watermark_video_simple(self, video_path, mask_coords, output_path=None):
        """
        简单视频去水印（逐帧处理）
        适合静态水印
        """
        if output_path is None:
            output_path = str(Path(video_path).with_suffix('')) + '_cleaned.mp4'
        
        # 读取视频
        clip = VideoFileClip(video_path)
        
        def process_frame(frame):
            # 转换颜色空间
            frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            
            # 创建mask
            mask = np.zeros(frame_bgr.shape[:2], dtype=np.uint8)
            x1, y1 = mask_coords[0]
            x2, y2 = mask_coords[1]
            mask[y1:y2, x1:x2] = 255
            
            # 修复
            result = cv2.inpaint(frame_bgr, mask, 3, cv2.INPAINT_TELEA)
            
            # 转换回RGB
            return cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
        
        # 处理视频
        processed_clip = clip.fl_image(process_frame)
        processed_clip.write_videofile(output_path, codec='libx264', audio_codec='aac')
        
        return output_path
    
    def remove_watermark_video_advanced(self, video_path, mask_coords_list, output_path=None):
        """
        高级视频去水印（支持动态水印）
        
        Args:
            video_path: 输入视频路径
            mask_coords_list: 每帧的水印坐标列表
            output_path: 输出路径
        """
        if output_path is None:
            output_path = str(Path(video_path).with_suffix('')) + '_cleaned_advanced.mp4'
        
        clip = VideoFileClip(video_path)
        
        def process_frame_with_tracking(frame, t):
            frame_idx = int(t * clip.fps)
            
            # 获取当前帧的mask坐标
            if frame_idx < len(mask_coords_list):
                mask_coords = mask_coords_list[frame_idx]
            else:
                mask_coords = mask_coords_list[-1]
            
            # 转换和处理
            frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            mask = np.zeros(frame_bgr.shape[:2], dtype=np.uint8)
            
            for coords in mask_coords:
                x1, y1 = coords[0]
                x2, y2 = coords[1]
                mask[y1:y2, x1:x2] = 255
            
            result = cv2.inpaint(frame_bgr, mask, 3, cv2.INPAINT_TELEA)
            return cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
        
        processed_clip = clip.fl(process_frame_with_tracking)
        processed_clip.write_videofile(output_path, codec='libx264', audio_codec='aac')
        
        return output_path
    
    def batch_process_images(self, input_dir, watermark_template=None):
        """批量处理图片去水印"""
        input_path = Path(input_dir)
        output_dir = input_path / 'cleaned'
        output_dir.mkdir(exist_ok=True)
        
        processed_files = []
        for img_file in input_path.iterdir():
            if img_file.suffix.lower() in self.supported_image_formats:
                try:
                    # 自动检测水印区域
                    regions = self.detect_watermark_region(str(img_file))
                    
                    if regions:
                        # 使用第一个检测到的区域
                        output_file = output_dir / f"{img_file.stem}_cleaned{img_file.suffix}"
                        self.remove_watermark_image(str(img_file), regions[0], str(output_file))
                        processed_files.append(str(output_file))
                    
                except Exception as e:
                    print(f"处理 {img_file} 时出错: {e}")
        
        return processed_files

# 创建全局实例
watermark_remover = WatermarkRemover()
    
# 图片去水印示例
# result = watermark_remover.remove_watermark_image("test.jpg", [(100, 100), (300, 200)])
# 视频去水印示例
# result = watermark_remover.remove_watermark_video_simple("test.mp4", [(100, 100), (300, 200)])
